/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.fnexecution.data;



import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnDataGrpc;

import com.bff.gaia.unified.runners.fnexecution.FnService;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.fn.data.UnifiedFnDataBufferingOutboundObserver;

import com.bff.gaia.unified.sdk.fn.data.UnifiedFnDataGrpcMultiplexer;

import com.bff.gaia.unified.sdk.fn.data.UnifiedFnDataInboundObserver;

import com.bff.gaia.unified.sdk.fn.data.CloseableFnDataReceiver;

import com.bff.gaia.unified.sdk.fn.data.FnDataReceiver;

import com.bff.gaia.unified.sdk.fn.data.InboundDataClient;

import com.bff.gaia.unified.sdk.fn.data.LogicalEndpoint;

import com.bff.gaia.unified.sdk.fn.stream.OutboundObserverFactory;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.stub.StreamObserver;

import com.bff.gaia.unified.vendor.guava.com.google.common.util.concurrent.SettableFuture;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.util.Queue;

import java.util.concurrent.*;



/**

 * A {@link FnDataService} implemented via gRPC.

 *

 * <p>This service allows for multiple clients to transmit {@link UnifiedFnApi.Elements} messages.

 *

 * <p>This service transmits all outgoing {@link UnifiedFnApi.Elements} messages to the first client

 * that connects.

 */

public class GrpcDataService extends UnifiedFnDataGrpc.UnifiedFnDataImplBase

    implements FnService, FnDataService {

  private static final Logger LOG = LoggerFactory.getLogger(GrpcDataService.class);



  public static GrpcDataService create(

      ExecutorService executor, OutboundObserverFactory outboundObserverFactory) {

    return new GrpcDataService(executor, outboundObserverFactory);

  }



  private final SettableFuture<UnifiedFnDataGrpcMultiplexer> connectedClient;

  /**

   * A collection of multiplexers which are not used to send data. A handle to these multiplexers is

   * maintained in order to perform an orderly shutdown.

   *

   * <p>TODO: (BEAM-3811) Replace with some cancellable collection, to ensure that new clients of a

   * closed {@link GrpcDataService} are closed with that {@link GrpcDataService}.

   */

  private final Queue<UnifiedFnDataGrpcMultiplexer> additionalMultiplexers;



  private final ExecutorService executor;

  private final OutboundObserverFactory outboundObserverFactory;



  private GrpcDataService(

      ExecutorService executor, OutboundObserverFactory outboundObserverFactory) {

    this.connectedClient = SettableFuture.create();

    this.additionalMultiplexers = new LinkedBlockingQueue<>();

    this.executor = executor;

    this.outboundObserverFactory = outboundObserverFactory;

  }



  /** @deprecated This constructor is for migrating Dataflow purpose only. */

  @Deprecated

  public GrpcDataService() {

    this.connectedClient = null;

    this.additionalMultiplexers = null;

    this.executor = null;

    this.outboundObserverFactory = null;

  }



  @Override

  public StreamObserver<UnifiedFnApi.Elements> data(

      final StreamObserver<UnifiedFnApi.Elements> outboundElementObserver) {

    LOG.info("Unified Fn Data client connected.");

    UnifiedFnDataGrpcMultiplexer multiplexer =

        new UnifiedFnDataGrpcMultiplexer(

            null, outboundObserverFactory, inbound -> outboundElementObserver);

    // First client that connects completes this future.

    if (!connectedClient.set(multiplexer)) {

      additionalMultiplexers.offer(multiplexer);

    }

    try {

      // We specifically return the connected clients inbound observer so that all

      // incoming messages are sent to the single multiplexer instance.

      return connectedClient.get().getInboundObserver();

    } catch (InterruptedException | ExecutionException e) {

      throw new RuntimeException(e);

    }

  }



  @Override

  public void close() throws Exception {

    // Cancel anything blocking on a client connecting to this service. This doesn't shut down the

    // Multiplexer, but if there isn't any multiplexer it prevents callers blocking forever.

    connectedClient.cancel(true);

    // Close any other open connections

    for (UnifiedFnDataGrpcMultiplexer additional : additionalMultiplexers) {

      try {

        additional.close();

      } catch (Exception ignored) {

        // Shutdown remaining clients

      }

    }

    if (!connectedClient.isCancelled()) {

      connectedClient.get().close();

    }

  }



  @Override

  @SuppressWarnings("FutureReturnValueIgnored")

  public <T> InboundDataClient receive(

      final LogicalEndpoint inputLocation,

      Coder<WindowedValue<T>> coder,

      FnDataReceiver<WindowedValue<T>> listener) {

    LOG.debug(

        "Registering receiver for instruction {} and target {}",

        inputLocation.getInstructionId(),

        inputLocation.getTarget());

    final UnifiedFnDataInboundObserver<T> observer =

        UnifiedFnDataInboundObserver.forConsumer(coder, listener);

    if (connectedClient.isDone()) {

      try {

        connectedClient.get().registerConsumer(inputLocation, observer);

      } catch (InterruptedException e) {

        Thread.currentThread().interrupt();

        throw new RuntimeException(e);

      } catch (ExecutionException e) {

        throw new RuntimeException(e.getCause());

      }

    } else {

      executor.submit(

          () -> {

            try {

              connectedClient.get().registerConsumer(inputLocation, observer);

            } catch (InterruptedException e) {

              Thread.currentThread().interrupt();

              throw new RuntimeException(e);

            } catch (ExecutionException e) {

              throw new RuntimeException(e.getCause());

            }

          });

    }

    return observer;

  }



  @Override

  public <T> CloseableFnDataReceiver<WindowedValue<T>> send(

      LogicalEndpoint outputLocation, Coder<WindowedValue<T>> coder) {

    LOG.debug(

        "Creating sender for instruction {} and target {}",

        outputLocation.getInstructionId(),

        outputLocation.getTarget());

    try {

      return UnifiedFnDataBufferingOutboundObserver.forLocation(

          outputLocation, coder, connectedClient.get(3, TimeUnit.MINUTES).getOutboundObserver());

    } catch (InterruptedException e) {

      Thread.currentThread().interrupt();

      throw new RuntimeException(e);

    } catch (TimeoutException e) {

      throw new RuntimeException("No client connected within timeout", e);

    } catch (ExecutionException e) {

      throw new RuntimeException(e);

    }

  }

}